import numpy as np
import matplotlib.pyplot as plt
from disba import PhaseDispersion
from bayesbay.discretization import Voronoi1D
import h5py

def read_inversion_results(npz_file):
    """
    读取保存的反演结果
    """
    data = np.load(npz_file, allow_pickle=True)
    save_dict = data['arr_0'].item()
    return save_dict

def compute_dispersion_curves(thickness, vs, vp_vs=1.77, rho_vp=[0.32, 0.77]):
    """
    根据给定的模型参数计算频散曲线
    """
    # 确保厚度数组正确（最后一个为半空间）
    thickness = thickness.copy()
    if len(thickness) > 0:
        thickness[-1] = 0  # 最后一层为半空间
    
    vp = vs * vp_vs
    rho = rho_vp[0] * vp + rho_vp[1]
    
    # 计算瑞雷波相速度
    pd = PhaseDispersion(thickness, vp, vs, rho)
    
    # 根据观测数据的频率范围计算理论曲线
    modes = []
    f = np.linspace(20,1,100)
    periods = 1/f
    for n in range(10):
        try:
            # 计算n阶模式
            r = pd(periods, mode=n, wave="rayleigh")
            modes.append((r.period, r.velocity, f'R{n}'))
        except Exception as e:
            print(f"Error computing R0 mode: {e}")
    print(f"Computed {len(modes)} modes")
        
    print(modes[-1])    
    return modes

def plot_dispersion_curves(save_dict, fv_file,output_file='dispersion_comparison.png'):
    """
    绘制频散曲线对比图
    """
    # 获取平均的S波速度结构
    statistics_vs = save_dict['statistics_vs']
    interp_depths = save_dict['interp_depths']
    mean_vs = statistics_vs['mean']
    
    # 构建分层模型用于计算理论频散曲线
    # 使用固定的厚度，因为我们只关注平均的Vs结构
    n_layers = len(interp_depths)
    # 创建一个简化的模型用于计算理论曲线
    # 这里我们创建一个简单的分层模型，层厚度基于反演结果的深度
    depths = interp_depths
    thickness = np.diff(depths)
    thickness = np.append(thickness, 0)  # 最后一层为半空间
    

    # 计算理论频散曲线
    theoretical_modes = compute_dispersion_curves(thickness, mean_vs)
    
    # 创建图形
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # 绘制观测数据
    # R0观测数据

    # 读取FV数据
    with h5py.File(fv_file, 'r') as f:
        FV = f['FV'][:]
        f_vals = f['f'][:]
        vs_vals = f['vs'][:]
    # 创建图形
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # 在背景上绘制FV能量分布
    X,Y = np.meshgrid(f_vals, vs_vals/1000)
    ax1.pcolormesh(X,Y, FV.T, cmap='Greys')

    r0_periods_obs = save_dict['R0_periods']
    r0_vel_obs = save_dict['R0_observed']
    ax1.plot(1/r0_periods_obs, r0_vel_obs, 'ro', label='Observed R0', markersize=6, alpha=0.7)
    
    # 如果存在R1观测数据
    if 'R1_periods' in save_dict and 'R1_observed' in save_dict:
        r1_periods_obs = save_dict['R1_periods']
        r1_vel_obs = save_dict['R1_observed']
        ax1.plot(1/r1_periods_obs, r1_vel_obs, 'bs', label='Observed R1', markersize=6, alpha=0.7)
    
    # 绘制理论频散曲线
    colors = ['red', 'blue', 'green']
    linestyles = ['-', '--', '-.']
    for i, (period, velocity, mode_name) in enumerate(theoretical_modes):
        if i < len(colors):
            ax1.plot(1/period, velocity, color=colors[i], linestyle=linestyles[i], 
                     label=f'Predicted {mode_name}', linewidth=2)
        else:
            ax1.plot(1/period, velocity, color='black',   linestyle='-', 
                     label=None, linewidth=2)
    
    ax1.set_xlabel('Frequency (Hz)', fontsize=12)
    ax1.set_ylabel('Phase Velocity (km/s)', fontsize=12)
    ax1.set_title('Dispersion Curves Comparison', fontsize=14)
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # 绘制Vs结构
    ax2.plot(mean_vs, interp_depths, 'r-', linewidth=2, label='Mean Vs')
    ax2.set_xlabel('Vs (km/s)', fontsize=12)
    ax2.set_ylabel('Depth (km)', fontsize=12)
    ax2.set_title('Inverted S-wave Velocity Profile', fontsize=14)
    ax2.invert_yaxis()  # 深度增加方向向下
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.show()
    
    print(f"Dispersion curves comparison plot saved to {output_file}")

def main():
    # 读取反演结果
    npz_file = 'figures/6.traceSeq.figures/test.inv/P06x.D123.Tp.NS.F1.0.3.5.V100.1400.NORM.STACK.FOLDER/P06x.D123.Tp.NS.F1.0.3.5.V100.1400.NORM.STACK.FOLDER_inversion_results.npz'
    npz_file = 'figures/6.traceSeq.figures/test.inv/P34x.D123.Tp.NS.F1.0.3.5.V100.1400.NORM.STACK.FOLDER/P34x.D123.Tp.NS.F1.0.3.5.V100.1400.NORM.STACK.FOLDER_inversion_results.npz'
    fv_file = 'figures/6.traceSeq.figures/test/P34x.D123.Tp.NS.F1.0.20.0.V100.1400.NORM.STACK.FOLDER.h5'
    
    save_dict = read_inversion_results(npz_file)
    print("Successfully loaded inversion results")
    
    # 显示一些基本信息
    print(f"Vs profile depths: {save_dict['interp_depths'].shape}")
    print(f"Mean Vs values: {save_dict['statistics_vs']['mean'].shape}")
    
    
    # 绘制原始频散曲线对比
    plot_dispersion_curves(save_dict,fv_file)

if __name__ == "__main__":
    main()